import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import os
import glob

# ============== 配置区域 ==============
# 字体配置
FONT_CONFIG = {
    'title_size': 20,
    'label_size': 20,
    'tick_size': 20,
    'legend_size': 20,
    'weight': 'normal',  # 去掉加粗，改为normal
    'family': 'Arial'  # 可选: 'Arial', 'Times New Roman', 'Helvetica'
}

# 颜色配置（字典格式）
COLOR_CONFIG = {
    'DE': '#1f77b4',        # 蓝色
    'DOWN': '#2CA02C',      # 绿色
    'KUBE': '#9467BD',      # 紫色
    'PSO': '#E377C2',       # 粉色
    'PPCM': '#BCBD22',      # 黄绿色
    'DEP_SOFT': '#FF7F0E',  # 橙色
    'GA': '#D62728',        # 红色
    'MONKEY': '#8C564B',    # 棕色
    'PPO': '#7F7F7F',       # 灰色
    'GNN': '#66CCFF',       # 浅蓝色 - 专门为GNN设置
}

# 图表配置
CHART_CONFIG = {
    'figsize': (12, 8),
    'dpi': 300,
    'alpha': 0.8,
    'bar_width_ratio': 0.8,
    'y_padding_ratio': 0.1,
    'legend_ncol_max': 5,
    'legend_alpha': 0.8
}
# =====================================

def plot_bar_chart_from_csv(csv_path: str, output_dir: str):
    """
    从CSV文件读取数据并生成条形图
    
    Args:
        csv_path (str): CSV文件路径
        output_dir (str): 输出目录
    """
    # 读取CSV数据
    try:
        df = pd.read_csv(csv_path, index_col=0)
    except Exception as e:
        print(f"错误: 无法读取文件 '{csv_path}': {e}")
        return
    
    if df.empty:
        print(f"警告: 文件 '{csv_path}' 为空")
        return
    
    # 获取文件名（不含扩展名）
    filename = os.path.splitext(os.path.basename(csv_path))[0]
    
    # 判断图表类型
    if filename.endswith('_Task'):
        chart_type = 'Task'
        x_label = 'Number of Tasks'
    elif filename.endswith('_Node'):
        chart_type = 'Node'
        x_label = 'Number of Nodes'
    else:
        chart_type = 'Task'  # 默认
        x_label = 'Number of Tasks'
    
    # 提取图表名称
    plot_name = filename.replace('_Task', '').replace('_Node', '')
    
    # 设置图表
    fig, ax = plt.subplots(figsize=CHART_CONFIG['figsize'])
    
    # 设置字体
    plt.rcParams['font.family'] = FONT_CONFIG['family']
    
    # 设置边框 - 只显示x轴和y轴实线
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(True)    # 显示y轴实线
    ax.spines['bottom'].set_visible(True)  # 显示x轴实线
    
    # 获取数据
    x_labels = df.index.tolist()
    series_names = df.columns.tolist()
    
    # 设置条形图参数
    n_groups = len(x_labels)
    n_series = len(series_names)
    bar_width = CHART_CONFIG['bar_width_ratio'] / n_series  # 总宽度平分给各个系列
    x_pos = np.arange(n_groups)
    
    # 绘制条形图
    bars = []
    for i, series_name in enumerate(series_names):
        # 从颜色字典中获取颜色，如果没有则使用默认颜色
        if series_name in COLOR_CONFIG:
            color = COLOR_CONFIG[series_name]
        else:
            # 如果系列名不在字典中，使用默认颜色循环
            color_list = list(COLOR_CONFIG.values())
            color = color_list[i % len(color_list)]
            
        values = df[series_name].values
        offset = (i - (n_series - 1) / 2) * bar_width
        bar = ax.bar(x_pos + offset, values, bar_width, 
                    label=series_name, color=color, alpha=CHART_CONFIG['alpha'])
        bars.append(bar)
    
    # 设置X轴
    ax.set_xlabel(x_label, fontsize=FONT_CONFIG['label_size'], fontweight=FONT_CONFIG['weight'])
    ax.set_xticks(x_pos)
    ax.set_xticklabels(x_labels, fontsize=FONT_CONFIG['tick_size'])
    
    # 设置Y轴（自适应）
    y_min = df.min().min()
    y_max = df.max().max()
    y_range = y_max - y_min
    y_padding = y_range * CHART_CONFIG['y_padding_ratio']  # 边距比例
    
    ax.set_ylim(y_min - y_padding, y_max + y_padding)
    ax.set_ylabel(plot_name.replace("_", " ").title(), 
                 fontsize=FONT_CONFIG['label_size'], fontweight=FONT_CONFIG['weight'])
    
    # 设置坐标轴数字字体大小
    ax.tick_params(axis='both', which='major', labelsize=FONT_CONFIG['tick_size'])
    
    # 设置图例（放在图表正上方）
    legend = ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15), 
                      fontsize=FONT_CONFIG['legend_size'], 
                      ncol=min(len(series_names), CHART_CONFIG['legend_ncol_max']))  # 最多N列
    legend.get_frame().set_linewidth(0)  # 去掉边框
    legend.get_frame().set_facecolor('white')  # 设置背景色
    legend.get_frame().set_alpha(CHART_CONFIG['legend_alpha'])  # 设置透明度
    
    # 去掉网格线
    ax.grid(False)
    
    # 设置标题（放在图表下方）
    ax.set_xlabel(f'{x_label}\n\n{plot_name.replace("_", " ").title()}', 
                 fontsize=FONT_CONFIG['label_size'], fontweight=FONT_CONFIG['weight'])
    
    # 调整布局（为图例留出更多空间）
    plt.subplots_adjust(top=0.85, bottom=0.15)
    
    # 保存图片
    output_path = os.path.join(output_dir, f'{filename}.png')
    plt.savefig(output_path, dpi=CHART_CONFIG['dpi'], bbox_inches='tight', 
                facecolor='white', edgecolor='none')
    plt.close()
    
    print(f"图表已保存为: '{output_path}'")

def process_all_csv_files(data_dir: str, output_dir: str):
    """
    处理指定目录中的所有CSV文件
    
    Args:
        data_dir (str): 包含CSV文件的目录
        output_dir (str): 输出目录
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 查找所有CSV文件
    csv_files = glob.glob(os.path.join(data_dir, '*.csv'))
    
    if not csv_files:
        print(f"错误: 在目录 '{data_dir}' 中未找到任何CSV文件")
        return
    
    print(f"找到 {len(csv_files)} 个CSV文件，开始处理...")
    
    for csv_file in csv_files:
        print(f"\n处理文件: {os.path.basename(csv_file)}")
        plot_bar_chart_from_csv(csv_file, output_dir)
    
    print(f"\n所有图表已保存到目录: '{output_dir}'")

if __name__ == '__main__':
    
    # 配置参数
    DATA_DIR = "data_chart/data"  # CSV文件目录
    OUTPUT_DIR = "test_res"       # 输出目录
    
    print("--- 开始从CSV数据生成图表 ---")
    print(f"数据目录: {DATA_DIR}")
    print(f"输出目录: {OUTPUT_DIR}")
    print(f"字体配置: {FONT_CONFIG}")
    print(f"颜色配置: {COLOR_CONFIG}")
    print(f"图表配置: {CHART_CONFIG}")
    
    # 处理所有CSV文件
    process_all_csv_files(DATA_DIR, OUTPUT_DIR)
